import random
import re
import sys

import math
import numpy as np
import torch
from PyQt5 import QtGui
from PyQt5.QtWidgets import *
from PIL import Image
from matplotlib import pyplot as plt
from torchvision import utils as vutils

from cloud_recognition_module import BaiduAipSpeech
from painter_module import Generator
from recorder_module import MyRecorder


class Main(QWidget):
    def __init__(self):
        super(Main, self).__init__()

        my_seed = np.random.randint(100)
        random.seed(my_seed)
        np.random.seed(my_seed)
        torch.manual_seed(my_seed)
        torch.cuda.manual_seed(my_seed)

        # ================== 加载神经网络生成器（鸟） ==================
        # self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.device = 'cpu'
        print(self.device)
        self.netG_bird = Generator().to(self.device)
        load_epoch = 350
        path_checkpointG = "./checkpointG_{}_epoch_cgan_birds.pkl".format(load_epoch)
        checkpointG = torch.load(path_checkpointG)
        self.netG_bird.load_state_dict(checkpointG['model_state_dict'])

        # ================== 加载神经网络生成器（脸） ==================
        # self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.device = 'cpu'
        print(self.device)
        self.netG_face = Generator().to(self.device)
        load_epoch = 100
        path_checkpointG = "./checkpointG_{}_epoch_cgan_faces.pkl".format(load_epoch)
        checkpointG = torch.load(path_checkpointG)
        self.netG_face.load_state_dict(checkpointG['model_state_dict'])

        # ================== 加载神经网络生成器（钱） ==================
        # self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.device = 'cpu'
        print(self.device)
        self.netG_money = Generator().to(self.device)
        load_epoch = 400
        path_checkpointG = "./checkpointG_{}_epoch_cgan_money.pkl".format(load_epoch)
        checkpointG = torch.load(path_checkpointG)
        self.netG_money.load_state_dict(checkpointG['model_state_dict'])

        # ============ initiate modules
        self.netG = None
        self.my_recorder = MyRecorder()
        self.baidu_smart_speech = BaiduAipSpeech()

        # ============ self attributes
        self.received_text = ""
        self.pic_num = 81
        self.keyword = None  # "中文词语"
        self.condition = None  # "English vocabulary"

        # ================== 条件字典 =====================
        self.firstcheck_keywords = "[(禽类)鸟脸(头像)钱钞]"

        self.condition_dict = {"red": [2., 0., 0.],  # birds condition
                            "yellow": [0., 2., 0.],
                              "blue": [0., 0., 1.],
                               "men": [1., 0.],      # faces condition
                             "women": [0., 1.],
                                 "1": [1., 0.],      # money condition
                               "100": [0., 1.],
                             "alice": [1., 0.],      # waifu condition
                             "chino": [0., 1.],
                              }

        self.condition_dictionary = {"红色": "red",       # 鸟
                                     "黄色": "yellow",
                                     "蓝色": "blue",
                                     "男": "men",       # 脸
                                     "女": "women",
                                     "一元": "1",         # 钱
                                     "百元": "100",
                                     "金发": "alice",     # 头像
                                     "银发": "chino",
                                     }

        self.kindnum = None

        # ============ GUI design
        self.resize(1200, 950)  # 横向，纵向
        self.setWindowTitle("心语画笔")

        self.label = QLabel(self)
        self.label.move(180, 30)
        self.label.setFixedSize(900, 900)
        self.label.setStyleSheet("QLabel{background:white;\
                                                 color:rgb(0,100,250,250);\
                                                 font-size:30px;\
                                                 font-weight:bold;\
                                                 font-family:Times New Roman;}")
        self.label.setText("          心语画笔 准备就绪")

        self.btn0 = QPushButton(self)
        self.btn0.setText("打开")
        self.btn0.move(25, 250)
        self.btn0.clicked.connect(self.open_image)

        self.btn1 = QPushButton(self)
        self.btn1.setText("录音")
        self.btn1.move(25, 50)
        self.btn1.clicked.connect(self.record)

        self.btn2 = QPushButton(self)
        self.btn2.setText("识别")
        self.btn2.move(25, 100)
        self.btn2.clicked.connect(self.recognize)

        self.btn3 = QPushButton(self)
        self.btn3.setText("取词")
        self.btn3.move(25, 150)
        self.btn3.clicked.connect(self.extract)

        self.btn4 = QPushButton(self)
        self.btn4.setText("绘画")
        self.btn4.move(25, 200)
        self.btn4.clicked.connect(self.paint)

        print('ready')

    # ====== 启动本地录音 =======
    def record(self):
        self.my_recorder.start_recording()
        self.label.setText('录音完成！我已听到。')

    # ====== 上传并识别 ======
    def recognize(self):
        self.received_text = self.baidu_smart_speech.start_recognition()
        self.label.setText(self.received_text)
        print(self.received_text)

    # ====== 提取关键词 ======
    def extract(self):
        kind = re.findall(self.firstcheck_keywords, self.received_text)
        print(kind)
        if "鸟" in kind or "禽类" in kind:

            # ================== 提取条件（鸟的） ====================
            self.netG = self.netG_bird
            self.kindnum = 3
            self.pic_num = 81
            red = re.findall("红色", self.received_text)
            yellow = re.findall("黄色", self.received_text)
            blue = re.findall("蓝色", self.received_text)
            if len(red) > 0:
                self.keyword = red[0]
                self.label.setText(self.keyword)

            elif len(yellow) > 0:
                self.keyword = yellow[0]
                self.label.setText(self.keyword)

            elif len(blue) > 0:
                self.keyword = blue[0]
                self.label.setText(self.keyword)

            else:
                self.keyword = "蓝色"
                self.label.setText('你说的颜色我以前没有见过...\n我给你画个蓝色的吧！^_^')

            print(self.keyword)
            self.condition = self.condition_dictionary[self.keyword]

        if "脸" in kind:

            # ================== 提取条件（脸的） ====================
            self.netG = self.netG_face
            self.kindnum = 2
            self.pic_num = 9
            
            if '男' in self.received_text:
                self.keyword = "男"
                self.label.setText('男')

            elif '女' in self.received_text:
                self.keyword = "女"
                self.label.setText('女')

            else:
                self.keyword = "女"
                self.label.setText('别纠结了，我还是给你画个女生吧！^_^')

            print(self.keyword)
            self.condition = self.condition_dictionary[self.keyword]

        if "钱" in kind or "钞" in kind:

            # ================== 提取条件（钱的） ====================
            self.netG = self.netG_money
            self.kindnum = 2
            self.pic_num = 81
            one = re.findall("一元", self.received_text)
            hundred = re.findall("百元", self.received_text)
            if len(hundred) > 0:
                self.keyword = hundred[0]
                self.label.setText(self.keyword)

            elif len(one) > 0:
                self.keyword = one[0]
                self.label.setText(self.keyword)

            else:
                self.keyword = "百元"
                self.label.setText('我以前学画的时候没有见过这种面值的...\n不如画百元大钞吧！^_^')

            print(self.keyword)
            self.condition = self.condition_dictionary[self.keyword]


    # ========== 条件生成图像 ================
    def paint(self):

        num = self.pic_num
        input_condition = torch.zeros((num, self.kindnum), dtype=torch.float, device=self.device)
        # print(input_condition)
        input_condition[:, ...] = torch.tensor(self.condition_dict[self.condition], device=self.device)
        # print(input_condition)
        z = torch.randn((num, 100 - self.kindnum, 1, 1), device=self.device)
        # print(z)
        c = input_condition.unsqueeze(2).unsqueeze(2).to(self.device)
        # print(c)
        x = torch.cat([z, c], dim=1)
        # print(x)
        generated_img = self.netG(x)
        # print(generated_img)
        img_grid = vutils.make_grid(generated_img.detach().cpu(), nrow=math.ceil(math.sqrt(num)), padding=2,
                                    normalize=True, scale_each=True)
        # print(img_grid)
        images = np.transpose(img_grid, (1, 2, 0))
        # print(images)
        plt.figure(figsize=(10, 10))
        plt.imshow(images)
        plt.axis('off')
        plt.savefig("心语画笔绘图结果.svg")
        print('图我已经帮你画好啦！打开文件夹看看吧^_^')
        self.label.setText('图我已经帮你画好啦！\n打开文件夹看看吧^_^')

    def open_image(self):
        img_name, _ = QFileDialog.getOpenFileName(self, caption="打开图片",
                                                        directory=".",
                                                        filter="*.png;;*.jpg;;All Files(*)")

        img = QtGui.QPixmap(img_name).scaled(self.label.width(), self.label.height())
        self.label.setPixmap(img)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    my_window = Main()
    my_window.show()
    sys.exit(app.exec())
